package com.mini.grpc.stream;

import com.mini.grpc.common.MethodDescriptor;
import com.mini.grpc.common.Metadata;
import com.mini.grpc.common.Status;
import com.mini.grpc.server.ServerCall;
import com.mini.grpc.server.ServerCallHandler;

import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 流式调用处理器
 * 处理各种类型的流式RPC调用
 * 
 * @param <ReqT> 请求类型
 * @param <RespT> 响应类型
 * @author Mini-gRPC
 */
public abstract class StreamingCallHandler<ReqT, RespT> implements ServerCallHandler<ReqT, RespT> {
    
    private final MethodDescriptor<ReqT, RespT> method;
    
    /**
     * 构造函数
     * 
     * @param method 方法描述符
     */
    protected StreamingCallHandler(MethodDescriptor<ReqT, RespT> method) {
        this.method = method;
    }
    
    @Override
    public ServerCall.Listener<ReqT> startCall(ServerCall<ReqT, RespT> call, Metadata headers) {
        StreamObserver<RespT> responseObserver = new ServerCallStreamObserver<>(call);
        
        switch (method.getType()) {
            case UNARY:
                return new UnaryServerCallListener(call, responseObserver);
            case CLIENT_STREAMING:
                return new ClientStreamingServerCallListener(call, responseObserver);
            case SERVER_STREAMING:
                return new ServerStreamingServerCallListener(call, responseObserver);
            case BIDI_STREAMING:
                return new BidiStreamingServerCallListener(call, responseObserver);
            default:
                throw new IllegalArgumentException("Unknown method type: " + method.getType());
        }
    }
    
    /**
     * 处理一元调用
     * 
     * @param request 请求
     * @param responseObserver 响应观察者
     */
    protected abstract void handleUnaryCall(ReqT request, StreamObserver<RespT> responseObserver);
    
    /**
     * 处理客户端流式调用
     * 
     * @param responseObserver 响应观察者
     * @return 请求观察者
     */
    protected abstract StreamObserver<ReqT> handleClientStreamingCall(StreamObserver<RespT> responseObserver);
    
    /**
     * 处理服务端流式调用
     * 
     * @param request 请求
     * @param responseObserver 响应观察者
     */
    protected abstract void handleServerStreamingCall(ReqT request, StreamObserver<RespT> responseObserver);
    
    /**
     * 处理双向流式调用
     * 
     * @param responseObserver 响应观察者
     * @return 请求观察者
     */
    protected abstract StreamObserver<ReqT> handleBidiStreamingCall(StreamObserver<RespT> responseObserver);
    
    /**
     * 服务端调用流观察者实现
     */
    private static class ServerCallStreamObserver<T> implements StreamObserver<T> {
        private final ServerCall<?, T> call;
        private final AtomicBoolean completed = new AtomicBoolean(false);
        
        ServerCallStreamObserver(ServerCall<?, T> call) {
            this.call = call;
        }
        
        @Override
        public void onNext(T value) {
            if (!completed.get()) {
                call.sendMessage(value);
            }
        }
        
        @Override
        public void onError(Throwable t) {
            if (completed.compareAndSet(false, true)) {
                Status status = Status.fromThrowable(t);
                call.close(status, null);
            }
        }
        
        @Override
        public void onCompleted() {
            if (completed.compareAndSet(false, true)) {
                call.close(Status.OK, null);
            }
        }
    }
    
    /**
     * 一元调用监听器
     */
    private class UnaryServerCallListener extends ServerCall.Listener<ReqT> {
        private final ServerCall<ReqT, RespT> call;
        private final StreamObserver<RespT> responseObserver;
        private ReqT request;
        
        UnaryServerCallListener(ServerCall<ReqT, RespT> call, StreamObserver<RespT> responseObserver) {
            this.call = call;
            this.responseObserver = responseObserver;
        }
        
        @Override
        public void onMessage(ReqT message) {
            if (request != null) {
                call.close(Status.INVALID_ARGUMENT.withDescription("Too many requests"), null);
                return;
            }
            request = message;
        }
        
        @Override
        public void onHalfClose() {
            if (request == null) {
                call.close(Status.INVALID_ARGUMENT.withDescription("No request received"), null);
                return;
            }
            
            try {
                handleUnaryCall(request, responseObserver);
            } catch (Exception e) {
                responseObserver.onError(e);
            }
        }
        
        @Override
        public void onCancel() {
            // 处理取消逻辑
        }
    }
    
    /**
     * 客户端流式调用监听器
     */
    private class ClientStreamingServerCallListener extends ServerCall.Listener<ReqT> {
        private final ServerCall<ReqT, RespT> call;
        private final StreamObserver<RespT> responseObserver;
        private final StreamObserver<ReqT> requestObserver;
        
        ClientStreamingServerCallListener(ServerCall<ReqT, RespT> call, StreamObserver<RespT> responseObserver) {
            this.call = call;
            this.responseObserver = responseObserver;
            this.requestObserver = handleClientStreamingCall(responseObserver);
        }
        
        @Override
        public void onMessage(ReqT message) {
            try {
                requestObserver.onNext(message);
            } catch (Exception e) {
                requestObserver.onError(e);
            }
        }
        
        @Override
        public void onHalfClose() {
            try {
                requestObserver.onCompleted();
            } catch (Exception e) {
                requestObserver.onError(e);
            }
        }
        
        @Override
        public void onCancel() {
            requestObserver.onError(new RuntimeException("Call cancelled"));
        }
    }
    
    /**
     * 服务端流式调用监听器
     */
    private class ServerStreamingServerCallListener extends ServerCall.Listener<ReqT> {
        private final ServerCall<ReqT, RespT> call;
        private final StreamObserver<RespT> responseObserver;
        private ReqT request;
        
        ServerStreamingServerCallListener(ServerCall<ReqT, RespT> call, StreamObserver<RespT> responseObserver) {
            this.call = call;
            this.responseObserver = responseObserver;
        }
        
        @Override
        public void onMessage(ReqT message) {
            if (request != null) {
                call.close(Status.INVALID_ARGUMENT.withDescription("Too many requests"), null);
                return;
            }
            request = message;
        }
        
        @Override
        public void onHalfClose() {
            if (request == null) {
                call.close(Status.INVALID_ARGUMENT.withDescription("No request received"), null);
                return;
            }
            
            try {
                handleServerStreamingCall(request, responseObserver);
            } catch (Exception e) {
                responseObserver.onError(e);
            }
        }
        
        @Override
        public void onCancel() {
            // 处理取消逻辑
        }
    }
    
    /**
     * 双向流式调用监听器
     */
    private class BidiStreamingServerCallListener extends ServerCall.Listener<ReqT> {
        private final ServerCall<ReqT, RespT> call;
        private final StreamObserver<RespT> responseObserver;
        private final StreamObserver<ReqT> requestObserver;
        
        BidiStreamingServerCallListener(ServerCall<ReqT, RespT> call, StreamObserver<RespT> responseObserver) {
            this.call = call;
            this.responseObserver = responseObserver;
            this.requestObserver = handleBidiStreamingCall(responseObserver);
        }
        
        @Override
        public void onMessage(ReqT message) {
            try {
                requestObserver.onNext(message);
            } catch (Exception e) {
                requestObserver.onError(e);
            }
        }
        
        @Override
        public void onHalfClose() {
            try {
                requestObserver.onCompleted();
            } catch (Exception e) {
                requestObserver.onError(e);
            }
        }
        
        @Override
        public void onCancel() {
            requestObserver.onError(new RuntimeException("Call cancelled"));
        }
    }
} 